package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

import nlp.langmodel.LanguageModel;
import nlp.util.Counter;
import nlp.util.CounterMap;

public class KatzBigramModel implements LanguageModel {
	private Collection<List<String>> trainingCollection;
	private Collection<List<String>> validationCollection;
	private EmpiricalUnigramLanguageModel unigramModel;
	private CounterMap<String, String> bigramCounterMap = new CounterMap<String, String>();
	private Counter<Integer> nCounter = new Counter<Integer>();
	private Counter<String> alphaCounter = new Counter<String>();
	
	private static final double INCREMENT_FACTOR = 1.0;
	private static final int THRESHOLD = 5;
	private static final String UNKNOWN = "<OOV>";
	private Double[] D = new Double[THRESHOLD + 1];
	private Double[] N = new Double[THRESHOLD + 2];
	private static final String START = "<S>";
	private static final String STOP = "</S>";
//	private double vocabulary = 0.0;
//	private double unseenCount = 0.0;
//	private double seenCount = 0.0;
	
	public KatzBigramModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.trainingCollection = trainingCollection;
		this.validationCollection = validationCollection;
		unigramModel = new EmpiricalUnigramLanguageModel(trainingCollection);
		train();
		smooth();
		calAlpha();
	}
	
	void train() {
		for (List<String> sentence : trainingCollection) {
			List<String> stopSentence = new ArrayList<String>(sentence);
			stopSentence.add(0, START);
			stopSentence.add(STOP);
			
			Iterator<String> iter = stopSentence.iterator();
			String first = iter.next();
			String second = iter.next();

			while (iter.hasNext() || !first.equals(STOP)) {
				Counter<String> counter = bigramCounterMap.getCounter(first);
				counter.incrementCount(second, INCREMENT_FACTOR);
				first = second;
				if (iter.hasNext()) {
					second = iter.next();
				}
			}
		}
		
		// find count of counts
		for (String first : bigramCounterMap.keySet()) {
			Counter<String> counter = bigramCounterMap.getCounter(first);
			for (String second : counter.keySet()) {
				double count = counter.getCount(second);
				nCounter.incrementCount((int)count, INCREMENT_FACTOR);
			}
		}
	}
	
	private void smooth() {
		Set<Integer> treeSet = new TreeSet<Integer>();
		treeSet.addAll(nCounter.keySet());
		Iterator<Integer> iter = treeSet.iterator();
		
		int next = 0;
		while (iter.hasNext() && next <= THRESHOLD) {
			next = iter.next();
			N[next] = nCounter.getCount(next);
			System.out.println("N[" + next + "] is " + N[next]);
		}
		N[0] = 0.0;
		
		for (int i = 1; i <= THRESHOLD; i++) {
			D[i] = (((i + 1) * N[i + 1]) / (i * N[i]) - ((THRESHOLD + 1) * N[THRESHOLD + 1]) / N[1])
					/
					(1.0 - ((THRESHOLD + 1) * N[THRESHOLD + 1]) / N[1]);
			System.out.println("D[" + i + "] is " + D[i]);
		}
		D[0] = 0.0;
		
	}
	
	private void calAlpha() {
		// for <OOV>, alpha is 1.0
		alphaCounter.setCount(UNKNOWN, INCREMENT_FACTOR);
		for (String history : bigramCounterMap.keySet()) {
			double sum = 0.0;
			double pUnigramSum = 0.0;
			Counter<String> counter = bigramCounterMap.getCounter(history);
			for (String word : counter.keySet()) {
				double count = counter.getCount(word);
				if ((int)count <= THRESHOLD) {
					sum += D[(int)count] * count;
				} else {
					sum += count;
				}
				pUnigramSum += unigramModel.getPWord(word);
			}
			double total = counter.totalCount();
			double alpha = ((total - sum) / total) / (1.0 - pUnigramSum);
			alphaCounter.setCount(history, alpha);
		}
	}
	
	double getWordLogProbability(List<String> sentence, int index) {
		if (index <= 0) {
			try {
				throw new Exception();
			} catch(Exception e) {
				System.err.println("Index out of range!");
			}
		}
		String history = sentence.get(index - 1);
		String word = sentence.get(index);
		return getPWord(history, word);
	}

	double getPWord(String history, String word) {
		if (!bigramCounterMap.containsKey(history)) {
			double alpha = alphaCounter.getCount(UNKNOWN);
			double pUni = unigramModel.getPWord(word);
//			System.out.println("word probability is " + Math.log(alpha * pUni) / Math.log(2.0) + ", alpha is " + alpha + ", and pUni is " + pUni);
			return Math.log(alpha * pUni) / Math.log(2.0);
		}
		
		Counter<String> counter = bigramCounterMap.getCounter(history);
		
		if (!counter.containsKey(word)) {
			double alpha = alphaCounter.getCount(history);
			double pUni = unigramModel.getPWord(word);
//			System.out.println("word probability is " + Math.log(alpha * pUni) / Math.log(2.0) + ", alpha is " + alpha + ", and pUni is " + pUni);
			
			if (Double.isInfinite(Math.log(alpha * pUni) / Math.log(2.0)) || Double.isNaN(Math.log(alpha * pUni) / Math.log(2.0))) {
//				System.out.println("word p is " + Math.log(alpha * pUni) / Math.log(2.0) + ", and the word is " + word);
//				System.out.println("and alpha is " + alpha + ", and pUni is " + pUni);
				// I'm sorry, I don't know what to do...
				return 0.0;
			}
			return Math.log(alpha * pUni) / Math.log(2.0);
		}
		
		double count = 0.0;
		int intCount = (int) counter.getCount(word);
		
		if (intCount <= THRESHOLD) {
			count = D[intCount] * intCount;
		} else {
			count = intCount;
		}
		double total = counter.totalCount();
//		System.out.println("word probability is " + Math.log(count / total) / Math.log(2.0) + ", count is " + count);
		return Math.log(count / total) / Math.log(2.0);
	}
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		double probability = 0.0;
		for (int i = 1; i < stoppedSentence.size(); i++) {
			double p = getWordLogProbability(stoppedSentence, i);
			probability += p;//getWordLogProbability(stoppedSentence, i);
		}
//		System.out.println("sentence p is " + probability);
		return Math.pow(2.0, probability);
	}

	String generateWord(String history) {
		Counter<String> counter = unigramModel.wordCounter;
		double sample = Math.random();
		double sum = 0.0;
		
		for (String word : counter.keySet()) {
			sum += Math.pow(2.0, getPWord(history, word));
			if (sum > sample) {
				return word;
			}
		}
		System.out.println("sum is " + sum);
//		return "*UNKNOWN*";
		return STOP;
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		
		String word = generateWord(START);
		while (!word.equals(STOP)) {
			sentence.add(word);
			word = generateWord(word);
		}
		return sentence;
	}

}
